#创建有key的二叉搜索树

#TreeNode:
#BinarySearchTree:

class TreeNode:
    '''二叉树'''
    def __init__(self, key, val, left=None, right=None, parent=None):
        self.key = key
        self.val = val
        self.left = left
        self.right = right
        self.parent = parent #显式记录父节点
        self.succ = self.findsuccessor() #后继节点
    def isLeftChild(self):
        return self.parent and (self.parent.left == self)
    def isRightChild(self):
        return self.parent and (self.parent.right == self)
    def isRoot(self):
        return not self.parent
    def isLeaf(self):
        return not (self.right or self.left)
    def hasBothChildren(self):
        return self.right and self.left
    def replaceNodeData(self, key, value, l, r): #替换节点
        self.key = key
        self.val = value
        self.left = l
        self.right = r
        if self.left: self.left.parent = self
        elif self.right: self.right.parent = self
        self.succ = self.findsuccessor()
    def __iter__(self): #迭代
        if self:
            if self.left:
                for elem in self.left:
                    yield elem
            yield self.key
            if self.right:
                for elem in self.right:
                    yield elem
    def findsuccessor(self): #寻找后继节点
        succ = None
        if self.right: #若有右子节点：succ=右子树左叶子
            succ = self.right.findMin()
        else:           #没有右子节点
            if self.parent:
                if self.isLeftChild(): #self为左子节点：succ=父节点
                    succ = self.parent 
                else:               #self为右子节点：succ=下一枝节点
                    self.parent.right = None
                    succ = self.parent.findsuccessor()
                    self.parent.right = self
        return succ
    def findMin(self): #左叶子
        current = self
        while current.left:
            current = current.left
        return current
    def spliceOut(self): #删除节点
        if self.isLeaf(): #若果是叶子节点：直接删
            if self.isLeftChild():
                self.parent.left = None
            else:
                self.parent.right = None
        else: #若果不是叶子节点
            if self.left: #左节点上移，删除节点和右子树
                if self.isLeftChild():
                    self.parent.left = self.left
                else:
                    self.parent.right = self.left
                self.left.parent = self.parent
            else:   #右节点上移，删除节点和左子树
                if self.isLeftChild():
                    self.parent.left = self.right
                else:
                    self.parent.right = self.right
                self.right.parent = self.parent
    # def __str__(self): #中序遍历打印
    #     result = []
    #     def inorder(tree):
    #         if tree:
    #             inorder(tree.left)
    #             result.append(tree.val)
    #             inorder(tree.right)
    #     inorder(self)
    #     return 'tree:'+str(result)

# a=TreeNode('3',3)
# b=TreeNode('5',5)
# c=TreeNode('4',4)
# d=TreeNode('6',6)
# tree = TreeNode('0',0)
# l=TreeNode('1',1,a,b,tree)
# r=TreeNode('2',2,c,d,tree)
# tree.left=l
# tree.right=r
# a.parent=l 
# b.parent=l 
# c.parent=r 
# d.parent=r

# #tree:[3, 1, 5, 0, 4, 2, 6]
# tree.spliceOut()
# print(tree)


class BinarySearchTree():
    '''二叉搜索树'''
    def __init__(self):
        self.root = None
        self.size = 0
    def __len__(self): #len()
        return self.size
    def __iter__(self): #迭代
        return self.root.__iter__()
    def __setitem__(self, key, val): #set()方法
        check = False
        if key in self: #如果key存在则替换
            check = True
        if self.root:
            self._put(key, val, self.root)
        else:
            self.root = TreeNode(key, val)
        if not check: #如果key不存在则插入
            self.size += 1
    def _put(self, key, val, curr): #递归查找put位置
        if key < curr.key: #如果key不存在则插入
            if curr.left:
                self._put(key, val, curr.left)
            else:
                curr.left = TreeNode(key, val, parent=curr)
        elif key == curr.key: #如果key存在则替换
            curr.val = val
        else:
            if curr.right:
                self._put(key, val, curr.right)
            else:
                curr.right = TreeNode(key, val, parent=curr)
    def __getitem__(self, key): #get()方法
        if self.root:
            res = self._get(key, self.root)
            if res:
                return res.val
        return None
    def _get(self, key, curr): 
        if not curr:
            return None
        elif curr.key == key:
            return curr
        elif key < curr.key:
            return self._get(key, curr.left)
        else:
            return self._get(key, curr.right)
    def __contains__(self, key): #in方法
        if self._get(key, self.root):
            return True
        else:
            return False
    def __delitem__(self, key): #del 方法
        if self.size > 1:
            nodeToRemove = self._get(key, self.root)
            if nodeToRemove:
                self.remove(nodeToRemove)
                self.size -= 1
            else:
                raise KeyError('Error, key not in tree')
        elif self.size == 1 and self.root.key == key:
            self.root = None
            self.size -= 1
        else:
            raise KeyError('Error, key not in tree')
    def remove(self, curr: TreeNode):
        if curr.isLeaf():
            if curr == curr.parent.left:
                curr.parent.left = None
            else: curr.parent.right = None
        elif curr.hasBothChildren():
            succ = curr.findsuccessor()
            succ.spliceOut()
            curr.key,curr.val = succ.key,succ.val
        else:
            if curr.left:
                if curr.isLeftChild():
                    curr.left.parent = curr.parent
                    curr.parent.left = curr.left
                elif curr.isRightChild():
                    curr.left.parent = curr.parent
                    curr.parent.right = curr.left
                else:
                    curr.replaceNodeData(curr.left.key,
                    curr.left.val, curr.left.left, curr.left.right)
            else:
                if curr.isLeftChild():
                    curr.right.parent = curr.parent
                    curr.parent.left = curr.right
                elif curr.isRightChild():
                    curr.right.parent = curr.parent
                    curr.parent.right = curr.right
                else:
                    curr.replaceNodeData(curr.right.key,
                    curr.right.val, curr.right.left, curr.right.right)

# mp = BinarySearchTree()
# for i in range(9):
#     mp[i] = i ** 2
# # for i in range(9):
# #     mp[i] = i ** 3
# del mp[4]
# print(mp.size)
# print(mp[8])
# print(mp.root)
